"""
.. module:: pymendel
   :platform: any
   :synopsis: This module models `Mendels Laws of Inheritance <https://en.wikipedia.org/wiki/Mendelian_inheritance#Mendel.27s_laws>`_.
    Rather than performing calculations, it actually allows you to perform
    crosses between individuals of a species, modeling the effects of random
    segregation and independent assortment.  So, if you want to see what might
    happen if you *actually* performed crosses to produce multiple offspring,
    you can use this module to see.

.. moduleauthor:: Pat Blair <pat@daburu.net>

"""

from random import randint
from sets import Set
# import threading
import uuid


class MendelianError(Exception):
    """Violations of `Mendel's Laws 
    <https://en.wikipedia.org/wiki/Mendelian_inheritance#Mendel.27s_laws>`_ or
    assumptions result in an exception of this type.
    """
    def __init__(self, value):
        self.value = value
    def __str__(self):
        return repr(self.value)
    
class Allele(object):
    """An allele is a discrete variation of a Gene.
    
    See also :py:class:`Gene`.
    """
    def __init__(self, dominant=False):
        """Don't call this constructor directly use 
        :py:meth:`Protogene.new_allele`.
        
        :param dominant: Is this the dominant allele?
        :type dominant: bool
        """
        # This is the trait observed in an individual when the allele is 
        # expressed (eg. 'tall' or 'dwarf').
        self._trait = None
        # This is the allele's symbol (eg. 'T' or 't'.
        self._symbol = None
        # Is this the dominant allele?
        self._dominant = dominant
        
    @property
    def dominant(self):
        """Is this the dominant allele?
        
        :returns: bool -- It is or it isn't.
        """
        return self._dominant

    @property
    def trait(self):
        """This is the trait observed in an individual when the allele 
        (*eg.* "tall" or "dwarf").
        
        :returns: str -- the phenotype
        """
        return self._trait
    
    @property
    def symbol(self):
        """This is the allele's symbol (*eg.* 'T' or 't').
        
        :returns: str -- the allele's symbol
        """
        return self._symbol
        
class Gene(object):
    """A gene is a unit of heredity in an organism.
    """
    def __init__(self, protogene, owner):
        """Don't call this constructor directly.  Create a 
        :py:class:`ProtoGene` first, then call :py:meth:`ProtoGene.new_gene`.  
        This is typically only a concern if you are developing the module or 
        extending its classes.  Most other times, you won't need to create new
        instances of this class directly.
                    
        :param owner: This is the individual that 'owns' this particular gene; 
            that is to say, the individual whose genotype is comprised of this 
            gene.
        :type dominant: :py:class:`Individual`
        :param protogene: This is the prototypical gene type that defines this
            gene.
        :type protogene: :py:class:`Protogene`
        """
        # We keep a reference to the individual that "owns" this gene (that is
        # to say, the individual whose genotype is constituted partially by
        # this gene).
        self._owner = owner
        # This is the gene's prototype.
        self._protogene = protogene
        # These are the alleles defined for this gene.
        self._alleles = [None, None]
    
    @property
    def alleles(self):
        """These are the valid alleles for this gene.
        
        :returns: list -- a list containing the two alleles for this gene
        """
        return self._alleles
            
    def _push_allele(self, symbol):
        """If you "push" an allele onto a gene in which neither of the alleles
        is defined, it becomes the first allele defined.  If you "push" another
        allele onto a gene, it becomes the second allele defined.  Beyond that,
        you'll raise a :py:exception::`MendelianError` because the base class
        assumes you're working with diploid individuals.
        
        :param symbol: the symbol of the allele
        :type symbol: str
        :raises: MendelianError -- If you attempt to use a symbol that is
                                   undefined for this gene, or you try to
                                   push to many alleles, you'll get an 
                                   exception.
        """
        # Get the allele for the symbol (or raise an exception if no such 
        # gene is defined.
        allele = None
        try:
            allele = self._protogene.get_allele(symbol)
        except KeyError:
            raise MendelianError(
                "The symbol %s is not defined for the %s gene." %
                symbol, self._protogene.name)
        # If the first allele hasn't been defined yet, set that one.
        if self._alleles[0] is None:
            self._alleles[0] = allele
            return
        # OK, the first allele has been set.  If the second allele slot is
        # empty, set that one.
        if self._alleles[1] is None:
            self._alleles[1] = allele
            # Make sure the dominant allele is first by reversing the list
            # if there is a dominant allele in the second position and a
            # recessive allele in the first.
            if self._alleles[1].dominant and not self._alleles[0].dominant:
                self._alleles.reverse()
            return
        raise MendelianError("The gene already has two alleles defined.")
    
    @property
    def _defined(self):
        """Has this gene been defined?  That is to say, have two alleles been
        supplied?
        
        :returns: bool -- *True* if the gene has been defined, otherwise 
                          *False*
        """
        return None not in self._alleles
                        
    @property
    def expression(self):
        """Get the dominant (*expressed*) allele for this gene.
        
        :returns: :py:class:`Allele` -- the dominant allele
        """
        # TODO: In this version, the API doesn't account for anything but simple dominance.
        # Return the first dominant allele.
        for a in self._alleles:
            if a.dominant:
                return a
        # We didn't return above?  We assume this means that no dominant 
        # alleles are present for this copy of the gene.  So, the recessive
        # allele is the one to express.  (In the current version of the module,
        # we assume that there is one dominant and one recessive allele.  When
        # it comes time to account for things like co-dominance, this would
        # be the place to do it.) 
        return self._alleles[0]
    
    @property
    def protogene(self):
        """This is the protogene upon which this gene is based.
        
        :returns: :py:class:`ProtoGene` -- this gene's prototype
        """
        return self._protogene
    
    @property
    def name(self):
        """
        This is a convenience property that returns the name of the protogene
        on which this copy of the gene is based.
        
        :See: :py:attr:`ProtoGene.name`
        """
        return self._protogene.name
    
    def _segregate(self):
        """
        The `Law of Segregation <https://en.wikipedia.org/wiki/Mendelian_inheritance#Law_of_Segregation_.28the_.22First_Law.22.29>`_
        states that a pair of alleles (assuming diploidy) is divided and one of 
        those alleles, selected randomly, is passed on to offspring. This 
        method returns the alleles for this copy of the gene in a random order.
        
        :returns: list of :py:class:`Allele` -- the alleles in random order
        """
        # We simulate the randomness of segregation during cell division by
        # returning the alleles in random order.
        i = randint(0,1)
        if i == 0:
            return [ self._alleles[0], self._alleles[1] ]
        else:
            return [ self._alleles[1], self._alleles[0] ]

class ProtoGene(object): 
    """A :py:class:`Protogene` is a prototype for the actual :py:class:`Gene`
    objects that individuals possess.
    """
    def __init__(self, name, dominant=False):
        """Create a new instance.
        
        :param name: This is the human-readable, friendly name for the gene.
                     Most of the time it should specify the characteristic
                     of an organism like "height" or "color".
        :type name: str
        :param dominant: Is this the dominant allele for this gene locus?
        :type dominant: bool -- *True* if it is, *False* if it isn't
        """
        self._name = name
        # This is an map of allele symbols to the alleles they symbolize.
        self._alleles_idx = {}
        

    def has_allele(self, symbol):
        """
        Does this gene have an allele with the symbol you provide?
        
        
        :returns: bool -- *True* if the gene contains an allele for the symbol,
                          otherwise *False*        
        """
        return symbol in self._alleles
    
    def get_allele(self, symbol):
        """Get the allele associated with a given symbol.
        
        :param symbol: This is the allele's symbol.
        :type symbol: str
        """
        return self._alleles_idx[symbol]
    
    @property
    def alleles(self):
        """Get the list of alleles defined for this gene.
        
        :returns: list of :py:class:`Allele`
        """
        return self._alleles_idx.values()
    
    @property
    def _symbols(self):
        """Get the list of symbols defined for this gene.
        
        :returns: list of str
        """
        return self._alleles_idx.keys()
    
    @property
    def name(self):
        """
        This is the friendly name by which the gene may be recognized.
        
        :returns: str
        """
        return self._name
    
    def new_allele(self, trait, symbol, dominant=False):
        """Define a new allele for this type gene.
        
        :param trait: See :py:attr:`Allele.trait`.
        :type trait: str
        :param symbol: See :py:attr:`Allele.symbol`.
        :type symbol: str
        :param dominant: See :py:attr:`Allele.dominant`.
        :type dominant: bool
        :returns: :py:class:`Allele`
        """
        # If this symbol has already been used to specify an allele...
        if symbol in self._alleles_idx:
            raise MendelianError("Symbol '%s' has already been used." % symbol)
        a = Allele(dominant)
        # Set the properties of the new allele.
        a._trait = trait
        a._symbol = symbol
        # The allele is 'of' this gene.
        a._of_gene  = self
        # If the new allele is marked as 'dominant', make sure that there 
        # isn't already a dominant allele specified.
        if dominant:
            for allele in self._alleles_idx.values():
                if allele.dominant:
                    raise MendelianError('Already specified a dominant allele')
        # Add the allele to the map, keyed to its symbol.
        self._alleles_idx[symbol] = a
        # Return the new allele to the caller.
        return a
        
    def _new_gene(self, owner):
        """Create a new :py:class:`Gene` instance based on this prototype.
        
        .. note::
        This is an internal method and you shouldn't need to use it for
        ordinary tasks.  Just create :py:class:`ProtoGene`s and
        :py:class:`Species` and the module will take care of this when you
        ask for new :py:class:`Individual`s.
        """
        # Create a new Gene instance with this object as its prototype.
        g = Gene(self, owner)
        return g
        

class Individual(object):
    """An individual is a single organism defined by its species.
    
    :See Also: :py:class:`Species`
    """    
    def __init__(self, species, name=None, genes=list()):
        """Under most circumstances, you don't want to call this constructor
        directly.  Instead, call :py:meth:`Species.new_individual` to create
        an individual defined by its species definition.
        
        :param species: This is the species to which the new individual
                        belongs.
        :type species: :py:class:`Species`
        """
        # This is the species to which the individual belongs.
        self._species = species
        # This is the individual's name (if any).
        self._name = name
        # This is the individual's ID.
        self._id = uuid.uuid4()
        # These are the genes that contstitute this individual's genotype.
        self._genes = list
        # This is a map of allele symbol (keys) to the genes (values) to which
        # they refer.  This index makes it faster for us set the phenotype from
        # a string of symbols. (See _update_indexes() and the genotype setter.)
        self._allele_idx = None
        # This is a map of genes indexed by name.  We use it when it comes
        # time to cross individuals and we must quickly find the homologues
        # between the two.
        self._gene_idx = None
        # The is the genotype, stored as a string of symbols.
        self._genotype_str = None

    @property
    def id(self):
        """This is an ID assigned to each new individual.
        
        :returns: UUID -- this individual's unique ID
        """
        return self._id

    @property 
    def genes(self):
        """These are the discrete genes that define this individual.
        
        :returns: list of :py:class:`Gene`
        :See Also: :py:attr:`Individual.genotype` property
        
        .. note:: 
            Before you call this, make sure you don't just need the *genotype*
            property.  If you just want the string of symbols, you can get it 
            there.
        
        """
        return self._genes
    
    @property
    def genotype(self):
        """This is the genotype of the individual represented as a string of
        allele symbols.
        """
        # If the genotype string was specifically set, just return it.
        if self._genotype_str is not None: 
            return self._genotype_str
        # Otherwise, we need to derive it.
        symbols = []
        for g in self._genes:
            for a in g._alleles:
                symbols.append(a.symbol)
        # Join the symbols to arrive at the string.
        self._genotype_str = ''.join(symbols)
        # Return the genotype string.
        return self._genotype_str
            
    
    @genotype.setter
    def genotype(self, s):
        """Set the genotype for this individual.
        
        :param s: This should be the string containing the allele symbols that
                  define the individual's genotype.
        :type s: str
        :raises: MendelianError -- If you attempt to set the genotype for an
                                   individual more than once, you'll get an 
                                   exception.
        """
        # TODO: Institute a lock here?
        # TODO: Split the string along any non-whitespace character.
        if self._genotype_str is not None:
            raise MendelianError("Individual's genotype has already been set.")
        # We assume each character to be a symbol for an allele.
        for c in s:
            # Get the gene for the allele to which the symbol refers.
            g = self._allele_idx[c]
            g._push_allele(c)
        # Make sure every gene has been defined.
        for g in self._genes:
            if not g._defined:
                raise MendelianError("One or more genes are not defined by " 
                                     "the genotype.")
        # Normalize the value by having the individual spit it back out, then
        # store it.
        self._genotype_str = self.genotype
    
    @property
    def phenotype(self):
        """An individual's `phenotype 
        <https://en.wikipedia.org/wiki/Phenotype>`_ is the composite of its
        observable traits.  This method returns a list of the traits
        specified by the dominant allele(s) for this gene.
        
        :See Also: :py:attr:`Allele.trait`
        :returns: list of str
        """
        return [g.expression.trait for g in self._genes]

        
    def _get_gene(self, name):
        """Get this indivudal's copy of a gene.
        
        :param name: This is the name of the gene.
        :type name: str
        :returns: :py:class:`Gene`
        :See Also: :py:attr:`Gene.name`
        """
        return self._gene_idx[name]
            
    @property
    def name(self):
        """This is the name of the individual. (You don't have to supply a 
        name unless it's useful.  If you don't, the individual's species name
        and an ID number are used to generate one.)
        
        :returns: str - the individual's name
        """
        # If this individual has been given a name, return it.
        if self._name is not None: return self._name
        # Create a name from the species and the ID.
        return ("%s %s" % (self._species.name, self._id)) 
    
    def get_trait(self, name):
        """Get the trait expressed by this individual in its phenotype.
        
        :param name: This is the name of the gene whose expression in this
                     individual you want to test.
        :type name: str
        :See Also: :py:attr:`Gene.name`
        :returns: str
        :See Also: :py:attr:`Allele.trait`
        :See Also: :py:attr:`Gene.expression`
        
        """
        # Get the gene.
        g = self._gene_idx[name]
        # Return the
        return g.expression.trait
            
    @property
    def traits(self):
        """Get the traits observed in this individual.
        
        :returns: {} - a map of gene names to the traits as they are expressed
                       in this individual
        """
        traits = {}
        for g in self._genes:
            traits[g.name] = self.get_trait(g.name)
        return traits
        
    def cross(self, other, name=None):
        """Cross this individual with another individual of the same species.
        
        :param other: This is the individual with whom you want to cross this
                      this individual.
        :type other: :py:class:`Individual`    
        :raises: MendelianError -- If you attempt to cross individuals of
                                   different species, you'll get an exception.    
        """
        # Make sure we're not trying to cross individuals of different species.
        if self._species._id != other._species._id:
            raise MendelianError("Cannot cross individuals of different"
                                 "species.")
        # Create the offpsring.
        offspring = self._species.new_individual(name=name)
        # Now we make a copy of each one of this individual's genes to pass
        # along with one allele from each parent.
        for g in self._genes:
            my_allele, _ = g._segregate()
            # Get the other indivudal's copy of the gene.
            other_g = other._get_gene(g.name)
            # Segregate the alleles of the other gene.
            other_allele, _ = other_g._segregate()
            # Get the offspring's copy of the gene.
            offspring_g = offspring._get_gene(g.name)
            # Push the alleles. (Recall that the method takes the allele's
            # symbol.)
            offspring_g._push_allele(my_allele.symbol)
            offspring_g._push_allele(other_allele.symbol)
        # The offspring now has a copy of each gene with alleles from the
        # parents.
        return offspring
             
    def is_identical(self, other):
        """
        Is this individual identical to another individual?
        
        :param other: This is the other individual.
        :type other: :py:class:`Individual`
        
        :returns: bool -- *True* if the individuals have identical genotypes.
        """
        # If either of the indivduals hasn't had its genotype set, there is
        # no basis of comparison for genetic identity.
        if self._genotype_str is None or other._genotype_str is None:
            return False
        # OK, so we have to genotypes to work with.  If they're the same, the
        # individuals are identical (at least so far as the traits we're
        # tracking are concerned.
        return self._genotype_str == other._genotype_str
        
    @property
    def species(self):
        """To what species does this individual belong?
        
        :returns: :py:class:`Species`
        """
        return self._species
    
    def _update_indexes(self):
        """:py:class:`Species` calls this when creating a new individual.
        It updates some indexes that make data-handling in the individual 
        object a little quicker.
        
        :See Also: :py:meth:`Species.new_individual`
        """
        # Create an index of all the possible alleles for each gene instance
        # so that we can more quickly get this individual's copy of the gene
        # given the allele.
        self._allele_idx = {}
        # We also want to be able to find the gene by name.
        self._gene_idx = {}
        # Let's do this.
        for g in self._genes:
            # Index this gene.
            self._gene_idx[g.name] = g
            # Now, index all of the alleles.
            for a in g.protogene.alleles:
                self._allele_idx[a.symbol] = g
    
class Species(object):
    """A species contains prototypical genes that define individual members
    of that species. 
    
    :See Also: :py:class:`Individual`
    """
    def __init__(self, name):
        """Create a new species.
        
        :param name: This is the name of your new species.
        :type name: str 
        """
        # This is the species ID.  (We use it for comparison.)
        self._id = uuid.uuid4()
        # This is the name of the species (ie. "mus musculus", "chicken", etc.).
        self._name = name
        # These are the gene prototypes that define the species.
        self._protogenes = []
        # We keep a flag to see if any individuals have been created yet.
        self._individuals_created = False

    @property
    def name(self):
        """This is the name of the species.
        
        :returns: str -- the name of the species
        """
        return self._name
    
    @property
    def protogenes(self):
        """
        Get a list of the genes that define this species.
        
        :returns: list of :py:class:`Protogene`
        """
        return self._protogenes
    
    def add_protogene(self, protogene):
        """Add a new gene prototype to this species.
        
        :param protogene: This is the protogene you want to add to the species.
        :type protogene: :py:class:`Protogene`
        
        :raises: MendelianError -- If you attempt to use a name for a protogene
                                   or allele symbol that has already been
                                   used for this species, you'll get an 
                                   exception.
        """
        # If any individuals have already been created, raise an error.  (You
        # can't redefine the species after some individuals are running
        # around.)
        if self._individuals_created:
            raise MendelianError("Cannot redefine the species after members "
                                 "have been created.")
        # Just for safety, do nothing if the caller sends in None...
        if protogene is None: return
        # ...or if we've already added it.
        if protogene in self._protogenes: return
        # Raise an exception if the caller is attempting to use a protogene 
        # name that has already been used.
        # (Yes, we're scanning the list, but it really shouldn't be that long.)
        for pg in self._protogenes:
            if pg.name == protogene.name:
                raise MendelianError("A protogene called '%s' is already "
                                     "defined for the species.")
            # Raise an exception of the caller is attempting to use a symbol
            # that has already been used.
            set1 = Set(pg._symbols)
            set2 = Set(protogene._symbols)
            # Get the intersection of the sets of symbols.
            isx = set1.intersection(set2)
            # If they have any of the same symbols, we have a problem.
            if len(isx) != 0:
                raise MendelianError("Detected an attempt to redefine one or "
                                     "more allele symbols.")
        # Add this one to the stack.
        self._protogenes.append(protogene)
        
    def new_individual(self, name=None):
        """Create a new individual of this species.
        
        :param name: See :py:class:`Individual.name`.
        :type name: str
        """
        # Create the new individual.  (Note that we're going to defer the
        # creation of the genes list so we have the individual in hand to
        # make the "owner" of each Gene as we create it.
        i = Individual(species=self, 
                       name=name,
                       genes=None)
        # OK, now let's create the genes from the protogenes defined for the
        # species.
        i._genes=[g._new_gene(i) for g in self._protogenes]
        # Update the new individual's indexes.
        i._update_indexes()
        # We have now created an individual, so the species definition is
        # locked.
        self._individuals_created = True
        return i
    
class Pmf(object):
    """This is a `probabilty mass function <https://en.wikipedia.org/wiki/Probability_mass_function>`_ 
    object for individuals.  You can populate it with a group of individuals,
    then query it for data structures that provide genotype, phenotype, and
    trait distributions.
    """
    def __init__(self):
        # This is a map of every individual added keyed by ID.
        self._individuals = {}
        # This is the map we use to track genotype frequencies.
        self._genotype_freq = {}
        # This is the map we use to track phenotype frequencies.
        self._phenotype_freq = {}
        # This is the map we use to track trait frequencies.
        self._trait_freq = {}
        # This is the species we're tracking.  (We'll set this when the first
        # individual is added.
        self._species_id = None
        # Here we have a flag that tells us whether or not the pmf has bee
        # initialized.
        self._initialized = False
        
    def add_individual(self, individual):
        """
        Add an individual to the population being tracked.
        
        :param individual: This is the individual to be added.
        :type individual: :py:class:`Individual`
        """
        # If the Pmf hasn't yet been initialized, do so now.
        if not self._initialized:
            self._initialize(individual)
            
        # Make sure the individual is of the same species as the initial
        # individual.  (This is how we make sure that only individuals of
        # one species are added to the population we're tracking.)
        if individual._species._id != self._species_id:
            raise MendelianError("The individual is not of the same species"
                                 "as the rest.")
        
        # If the individual has already been added, raise an exception.
        if individual.id in self._individuals:
            raise MendelianError("The individual has already been added.")
        # Make note of this individual's ID.
        self._individuals[individual.id] = individual
        #=======================================================================
        # Add the individual's genotype to the genotype frequency map and the
        # phenotype to the phenotype frequency map.
        #=======================================================================
        # If we haven't seen this genotype before...
        if not individual.genotype in self._genotype_freq:
            # ...add an entry for it in the genotype frequeny map.
            self._genotype_freq[individual.genotype] = 0.0
        # Get the individual's phenotype.
        individual_phenotype = ','.join(individual.phenotype)
        # If we haven't seen this phenotype before...
        if not individual_phenotype in self._phenotype_freq:
            # ...add an entry for it in the phenotype frequency map.
            self._phenotype_freq[individual_phenotype] = 0.0
        # Update the genotype and phenotype frequencies.
        self._genotype_freq[individual.genotype] += 1.0
        self._phenotype_freq[individual_phenotype] += 1.0
        # Add all of the individual's traits to the trait frequency map.
        for g in individual.genes:
            trait = individual.get_trait(g.name)
            self._trait_freq[g.name][trait] += 1.0
        
    def _initialize(self, individual):
        """We can speed things up by initialzing the maps when the first 
        individual is added.
        
        :param individual: This is the individual upon which we're going to
                           set up the internal indices.
        :type individual: :py:class:`Individual`
        """
        # This is just a sanity check.
        if self._initialized:
            return
        # Initialize the phenotype frequency map.
        for g in individual.genes:
            allele_freqs = {}
            for a in g.protogene.alleles:
                allele_freqs[a.trait] = 0.0
            self._trait_freq[g.name] = allele_freqs
        # Make note of the species of the initial individual.
        self._species_id = individual._species._id
        # Set the flag.  We're initialized.
        self._initialized = True
            
                
    def get_genotype_distribution(self):
        """Get the genotype distribution for the population.
        
        :returns: {} -- Each distinct genotype is a key and its relative 
                        frequency (*ie.* a float between 0 and 1 inclusive) is 
                        the value.
        """
        # Determine the total number of individuals.
        total = float(len(self._individuals))
        # Create a map to hold the distributions.
        m = {}
        for gtype in self._genotype_freq.keys():
            m[gtype] = (self._genotype_freq[gtype]/total)
        return m
    
    def get_phenotype_distribution(self):
        """Get the phenotype distribution for the population.
        
        :returns: {} -- Each distinct phenotype is a key and its relative
                        frequency (*ie.* a float between 0 and 1 inclusive) is 
                        the value.
        """
        # Determine the total number of individuals.
        total = float(len(self._individuals))
        # Create a map to hold the distributions.
        m = {}
        for ptype in self._phenotype_freq.keys():
            m[ptype] = (self._phenotype_freq[ptype]/total)
        return m
    
    def get_trait_distribution(self):
        """Get the trait distribution for the population.
        
        :returns: {} -- a data structure in which each gene's name is a key
                        and the values are dictionaries with each possible
                        allele of the gene and its relative frequency
        """
        # How many individuals are in the population?
        total = float(len(self._individuals))
        # Create the dictionary object.
        m = {}
        # We start with the trait frequency dictionary.  We're going to
        # go through each gene's name.
        for gname in self._trait_freq.keys():
            # Get the dictionary of trait frequencies for this gene.
            trait_freqs = self._trait_freq[gname]
            # We need a new dictionary to hold all of the trait frequencies.
            trait_dists = {}
            # Now determine the relative frequencies for each trait.
            for trait_name in trait_freqs.keys():
                freq = trait_freqs[trait_name]
                trait_dists[trait_name] = freq/total
            m[gname] = trait_dists
        return m
    

    
        
    
        
